package com.example.code.class03;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.Set;

/**
 * @Date 2022/7/17
 * @Created by Jonathan
 */
public class Code09_DistanceKNodes {
    public static void main(String[] args) {
        // test createParentMap
        Map<Node, Node> map = new HashMap<>();
        final Node root = createTree();
        createParentsMap(root, map);
        System.out.println(map);
        // test  getDistanceKNodes
        List<Node> kNodes = getDistanceKNodes(root, root.left, 2);
        System.out.println("ans:" + kNodes);

    }

    static class Node {
        Integer value;
        Node left;
        Node right;

        public Node(Integer value) {
            this.value = value;
        }

        @Override
        public String toString() {
            return "Node[" +
                    "" + value +
                    ']';
        }
    }


    private static void createParentsMap(Node cur, Map<Node, Node> parentMap) {
        // base case
        if (cur == null) {
            return;
        }
        // 如果当前节点的左孩子存在的话 左孩子的parent就是当前节点
        if (cur.left != null) {
            parentMap.put(cur.left, cur);
            createParentsMap(cur.left, parentMap);
        }

        // 如果当前节点的右孩子存在的话 右孩子的parent就是当前节点
        if (cur.right != null) {
            parentMap.put(cur.right, cur);
            createParentsMap(cur.right, parentMap);
        }
    }


    private static List<Node> getDistanceKNodes(Node root, Node target, int k) {
        Map<Node, Node> parentMap = new HashMap<>();
        createParentsMap(root, parentMap);
        // 记录访问过的node
        Set<Node> visited = new HashSet<>();
        List<Node> ans = new ArrayList<>();
        Queue<Node> queue = new LinkedList<>();
        queue.add(target);
        int currentLevel = 0;
        while (!queue.isEmpty()) {
            // 超过距离就结束循环
            if (currentLevel > k) {
                break;
            }
            int currentLevelSize = queue.size();
            while (currentLevelSize > 0) {
                Node cur = queue.poll();
                // 距离相等就加入ans
                if (currentLevel == k) {
                    ans.add(cur);
                }
                visited.add(cur);
                if (cur.left != null && !visited.contains(cur.left)) {
                    queue.add(cur.left);
                }
                if (cur.right != null && !visited.contains(cur.right)) {
                    queue.add(cur.right);
                }
                if (parentMap.get(cur) != null && !visited.contains(parentMap.get(cur))) {
                    queue.add(parentMap.get(cur));
                }
                currentLevelSize--;
            }
            // 层次加一
            currentLevel++;
        }
        return ans;
    }

    private static Node createTree() {
        Node root = new Node(1);
        Node node2 = new Node(2);
        Node node3 = new Node(3);
        Node node4 = new Node(4);
        Node node5 = new Node(5);
        Node node6 = new Node(6);
        Node node7 = new Node(7);
        root.left = node2;
        root.right = node3;
        node2.left = node4;
        node2.right = node5;
        node4.left = node6;
        node4.right = node7;
        return root;
    }


}
